#include <unistd.h>
#include <time.h>
#include <mutex>
#include <list>
#include <sys/syscall.h>
#include <string>
#include <iostream>
#include <fstream>
#include <sstream>

class DclLlmProfiler
{
public:
    DclLlmProfiler(const std::string& cat, const std::string& name, const std::string& args)
        : cat_(cat), name_(name), args_(args)
    {
        struct timespec ts;
        clock_gettime(CLOCK_REALTIME, &ts);
        ts_ = ts.tv_sec * 1000000000 + ts.tv_nsec;

        pid_ = getpid();
        tid_ = syscall(SYS_gettid);
    }

    ~DclLlmProfiler()
    {
        struct timespec ts;
        clock_gettime(CLOCK_REALTIME, &ts);
        int64_t end_ts = ts.tv_sec * 1000000000 + ts.tv_nsec;
        dur_ = end_ts - ts_;

        std::unique_lock<std::mutex> lock(s_profMtx);
        std::ostringstream oss;
        oss << *this;
        s_profList.emplace_back(std::move(oss).str());
    }

    friend std::ostream& operator<<(std::ostream& os, const DclLlmProfiler& prof);

    static std::mutex s_profMtx;
    static std::list<std::string> s_profList;

private:
    std::string cat_;    // category，例如cpu_op、nnp_op
    std::string name_;   // 自主定义的名称
    std::string args_;   // 要记录的参数

    int64_t ts_;         // 起始时间，单位为ns
    int64_t dur_;        // 持续时间，单位为ns
    int32_t pid_;
    int32_t tid_;
};

std::mutex DclLlmProfiler::s_profMtx;
std::list<std::string> DclLlmProfiler::s_profList;

#define DCL_LLM_RECORD_FUNCTION(cat, name, args) \
    DclLlmProfiler dclLlmProfiler(cat, name, args);

std::ostream& operator<<(std::ostream& os, const DclLlmProfiler& prof)
{
    // 这里先原样输出，最后由python脚本处理为json格式
    os << prof.cat_ << "\n"
        <<  prof.name_ << "\n"
        <<  prof.pid_ << "\n"
        <<  prof.tid_ << "\n"
        << prof.ts_ << "\n"
        << prof.dur_ << "\n"
        << prof.args_ << "\n";
    return os;
}


extern "C" void my_func2(int k)
{
    DCL_LLM_RECORD_FUNCTION("cpu_op", "my_func2", std::string("k=") + std::to_string(k));
    usleep(2000);
}

extern "C" void my_func(int x)
{
    DCL_LLM_RECORD_FUNCTION("nnp_op", "my_func", std::string("x=") + std::to_string(x));
    my_func2(x * 2);
    usleep(1000);
}

extern "C" void dclLlmStopProfiler(void)
{
    std::ofstream ofs("dcl_llm_profile.txt");
    for (const auto& prof : DclLlmProfiler::s_profList) {
        ofs << prof;
    }
}